import os
import re
import shutil
import random
from collections import defaultdict

# 路径设置
source_train_dir = os.path.join('data', 'NEU-DET', 'train', 'images')
target_train_dir = os.path.join('data', 'NEU-DET', 'train_reorganized')
target_test_dir = os.path.join('data', 'NEU-DET', 'test')

# 确保目标目录存在
os.makedirs(target_train_dir, exist_ok=True)
os.makedirs(target_test_dir, exist_ok=True)

# 创建good（无缺陷）类别的文件夹
os.makedirs(os.path.join(target_train_dir, 'good'), exist_ok=True)
os.makedirs(os.path.join(target_test_dir, 'good'), exist_ok=True)
print("已创建good（无缺陷）类别的文件夹")

# 获取所有图片文件并按缺陷类型分组
defect_files = defaultdict(list)

for file in os.listdir(source_train_dir):
    if file.endswith('.jpg') or file.endswith('.png'):
        match = re.match(r'([a-zA-Z_]+)_\d+', file)
        if match:
            defect_type = match.group(1)
            defect_files[defect_type].append(file)

# 为每种缺陷类型创建文件夹并分配图片
for defect_type, files in defect_files.items():
    # 创建训练集和测试集文件夹
    train_defect_dir = os.path.join(target_train_dir, defect_type)
    test_defect_dir = os.path.join(target_test_dir, defect_type)
    
    os.makedirs(train_defect_dir, exist_ok=True)
    os.makedirs(test_defect_dir, exist_ok=True)
    
    # 随机选择10张图片用于测试集
    random.seed(42)  # 设置随机种子以确保结果可重现
    test_files = random.sample(files, 10)
    
    # 复制文件到相应目录
    for file in files:
        source_path = os.path.join(source_train_dir, file)
        
        if file in test_files:
            # 复制到测试集
            target_path = os.path.join(test_defect_dir, file)
        else:
            # 复制到训练集
            target_path = os.path.join(train_defect_dir, file)
        
        shutil.copy2(source_path, target_path)
        print(f"已复制: {file} -> {'测试集' if file in test_files else '训练集'}")

print("\n数据集重组完成!")
print(f"训练集目录: {target_train_dir}")
print(f"测试集目录: {target_test_dir}")

# 打印统计信息
print("\n统计信息:")
for defect_type in sorted(list(defect_files.keys()) + ['good']):
    train_dir = os.path.join(target_train_dir, defect_type)
    test_dir = os.path.join(target_test_dir, defect_type)
    
    train_count = len(os.listdir(train_dir)) if os.path.exists(train_dir) else 0
    test_count = len(os.listdir(test_dir)) if os.path.exists(test_dir) else 0
    
    print(f"- {defect_type}: 训练集 {train_count}张, 测试集 {test_count}张") 